import os
import operator
from dotenv import load_dotenv
from typing import Annotated, Union
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_openai import ChatOpenAI

# --- Pydantic & LangGraph Imports ---
# 1. 从 Pydantic 导入 Field 和 BaseModel
from pydantic import Field, BaseModel
# 2. 从 langgraph.graph 导入核心组件
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

# --- 1. 加载环境变量 ---
load_dotenv()

# --- 2. 定义工具 ---
search_tool = DuckDuckGoSearchRun()
tools = [search_tool]

# --- 3. 定义状态 (Pydantic 版本) ---
# 这是本次教程的核心升级点。
# 我们不再使用 TypedDict，而是定义一个继承自 langgraph.graph.BaseModel 的 Pydantic 模型。
# BaseModel 实际上是 Pydantic.BaseModel 的一个别名，但建议从 langgraph 导入以保证兼容性。
class AgentState(BaseModel):
    """
    使用 Pydantic 定义的图状态。
    这提供了自动的类型验证和更好的代码提示。
    """
    # a. 定义 messages 字段
    #    - 类型提示为 list[BaseMessage]，确保所有消息都是 LangChain 的标准消息对象。
    #    - 我们使用 `Annotated` 来告诉 LangGraph 如何处理这个字段的更新。
    #    - `operator.add` 是一个 "reducer"，它指定了新的消息应该被
    #      添加（append）到现有列表的末尾，而不是替换整个列表。
    #      这对于维护对话历史至关重要，无论状态是用 TypedDict 还是 Pydantic 定义。
    #    - Field(default_factory=list) 依然是 Pydantic 的推荐做法，用于设置默认值。
    messages: Annotated[list[BaseMessage], operator.add] = Field(default_factory=list)

# --- 4. 定义图中的节点 ---

# 初始化模型并绑定工具
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(tools)

# a. Agent 节点
def agent_node(state: AgentState):
    """
    调用 LLM 来决定下一步行动。
    注意：输入参数 `state` 现在是 AgentState 的一个实例，而不是字典。
    """
    print("---AGENT: 思考中...---")
    # 访问状态时，我们使用 `state.messages` 而不是 `state["messages"]`。
    # 这得益于 Pydantic，能提供更好的静态类型检查和自动补全。
    response = model_with_tools.invoke(state.messages)
    
    # 返回一个字典，LangGraph 会自动用它来更新状态对象。
    # Pydantic 会在后台验证这里的 "messages" 键和它的值是否符合 AgentState 的定义。
    return {"messages": [response]}

# b. Tool 节点
# ToolNode 的工作方式完全不变。它能智能地处理 Pydantic 状态或 TypedDict 状态。
tool_node = ToolNode(tools)

# --- 5. 定义条件边 ---
def should_continue(state: AgentState) -> str:
    """
    根据 Agent 的最新响应决定工作流的下一站。
    同样，这里也使用 `state.messages`。
    """
    print("---ROUTER: 决策中...---")
    last_message = state.messages[-1]

    if not last_message.tool_calls:
        print("---ROUTER: 任务完成，结束流程。---")
        return "__end__"
    else:
        print("---ROUTER: 需要调用工具。---")
        return "tools"

# --- 6. 构建图 ---
# 当使用 Pydantic 模型作为状态时，StateGraph 的初始化方式完全相同。
# LangGraph 会自动识别它是一个 Pydantic 模型。
workflow = StateGraph(AgentState)

# 添加节点、设置入口点、添加边的逻辑也和以前完全一样。
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
    "agent",
    should_continue,
)
workflow.add_edge("tools", "agent")

app = workflow.compile()

# --- 7. 运行图 ---
if __name__ == "__main__":
    question = "通义千问的最新版本是什么？"
    
    # 输入的格式依然是一个字典，LangGraph 会负责将其加载到 AgentState 模型中。
    inputs = {"messages": [HumanMessage(content=question)]}
    
    for event in app.stream(inputs, stream_mode="values"):
        # `event` 是一个字典，其键与 AgentState 的字段匹配。
        # 我们可以通过 `event['messages'][-1]` 来访问最新的消息。
        event['messages'][-1].pretty_print()
        print("\n" + "="*50 + "\n") 